import cv2
import numpy as np
import subprocess

from video_base import VideoBase


class VideoWatermarkRemoval(VideoBase):
    def __init__(self,
                 input_video_path: str,
                 output_name="video_watermark_removal.mp4",
                 roi: (int, int, int, int) = None):
        """
        input_video_path: 切割视频地址
        output_name: 输出视频名称
        roi: logo位置 x1,y1,x2,y2
        """
        self.input_video_path = input_video_path
        self.output_name = output_name
        # 设置水印区域
        if roi is None:
            self.logo_x1 = None
            self.logo_y1 = None
            self.logo_x2 = None
            self.logo_y2 = None
            self.logo_width = None
            self.logo_height = None
        else:
            self.logo_x1 = roi[0]
            self.logo_y1 = roi[1]
            self.logo_x2 = roi[2]
            self.logo_y2 = roi[3]
            self.logo_width = roi[2] - roi[0]
            self.logo_height = roi[3] - roi[1]
        super().__init__(self.input_video_path)
        # 创建 mask
        self.mask = np.zeros((self.height, self.width), dtype=np.uint8)
        if self.logo_x1:
            self.mask[self.logo_y1:self.logo_y2, self.logo_x1:self.logo_x2] = 255

    @staticmethod
    def _select_roi(frame):
        """
        在视频首帧手动选择水印区域。
        参数:
            frame (np.ndarray): 视频首帧图像
        返回:
            tuple: (x, y, w, h) 水印区域坐标
        """
        print("请用鼠标框选水印区域，按 ENTER 确认")
        roi = cv2.selectROI("选择水印区域", frame, False, False)
        cv2.destroyWindow("选择水印区域")
        return roi

    def _save_first_frame(self, frame, filename="first_frame.jpg"):
        cv2.rectangle(frame, (self.logo_x1, self.logo_y1), (self.logo_x2, self.logo_y2), (0, 255, 0), 2)
        cv2.imwrite(filename, frame)
        print(f"第一帧已保存为: {filename}")

    def _get_roi(self):
        self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        ret, frame = self.cap.read()
        if ret:
            if self.logo_x1 is None:
                roi = self._select_roi(frame)
                self.logo_x1, self.logo_y1, self.logo_width, self.logo_height = roi
                self.logo_x2 = self.logo_x1 + self.logo_width
                self.logo_y2 = self.logo_y1 + self.logo_height
                print(f"x1: {self.logo_x1},y1:{self.logo_y1},x2:{self.logo_x2},y2:{self.logo_y2}")
                self.mask[self.logo_y1:self.logo_y2, self.logo_x1:self.logo_x2] = 255
            return frame
        else:
            print("无法读取视频第一帧")

    def get_first_frame(self):
        """
        获取第一帧的图片信息
        :return:
        """
        frame = self._get_roi()
        if frame is not None:
            self._save_first_frame(frame)

    def remove_watermark_ffmpeg(self):
        """
        使用 FFmpeg delogo 滤镜去除视频水印
        :return: 去水印后的视频路径
        """
        cmd = [self.ffmpeg_exe, "-y"]  # 覆盖输出文件
        # 如果启用 GPU，加上 hwaccel 参数
        # if is_cuda_available():
        #     cmd += ["-hwaccel", "cuda"]
        # 输入视频
        cmd += ["-i", self.input_video_path]
        # delogo 滤镜
        cmd += [
            "-vf",
            f"delogo=x={self.logo_x1}:y={self.logo_y1}:w={self.logo_width}:h={self.logo_height}:show=0",
            "-c:a", "copy",  # 保留原音频
            self.output_name
        ]
        print("执行命令：", " ".join(cmd))
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if result.returncode == 0:
            print(f"水印去除完成，输出文件：{self.output_name}")
        else:
            print("去水印失败，错误信息：", result.stderr.decode())

    def remove_watermark(self, show_video=False):
        """
        视频水印去除
        :return:
        """
        fourcc = cv2.VideoWriter_fourcc(*"XVID")  # 用 AVI + XVID 避免黑屏
        writer = cv2.VideoWriter(self.output_name, fourcc, self.fps, (self.width, self.height))
        x1, y1, x2, y2 = self.logo_x1, self.logo_y1, self.logo_x2, self.logo_y2
        roi_mask = self.mask[y1:y2, x1:x2]

        while True:
            ret, frame = self.cap.read()
            if not ret:
                break
            orig = frame.copy()
            roi = frame[y1:y2, x1:x2]
            # if is_cuda_available():
            #     # GPU 版本：上传到 GpuMat（OpenCV CUDA）
            #     gpu_roi = cv2.cuda_GpuMat()
            #     gpu_roi.upload(roi)
            #     # GPU 没有 inpaint，目前仍用 CPU 修复
            #     repaired = cv2.inpaint(roi, roi_mask, 3, cv2.INPAINT_TELEA)
            #     # dst = cv2.inpaint(src, inpaintMask, inpaintRadius, flags)
            #     #   src: 原始图像（彩色或灰度）
            #     #   inpaintMask: 掩码（同 src 大小，单通道，非零像素表示要修复的区域）
            #     #   inpaintRadius: 修复半径（3~5 常用，越大修复范围越广，但可能模糊）
            #     #   flags: 修复算法类型
            #     #       cv2.INPAINT_TELEA → 基于 Telea 算法，速度快，效果自然
            #     #       cv2.INPAINT_NS → 基于 Navier-Stokes 方程，效果更平滑但慢一些
            #     #   返回值 dst 是修复后的图像，即去掉水印后的帧。
            # else:
            repaired = cv2.inpaint(roi, roi_mask, 3, cv2.INPAINT_TELEA)

            # 放回原帧
            frame[y1:y2, x1:x2] = repaired

            # 写入视频
            writer.write(frame)

            # ---------------------------
            # 2. 显示预览（可选）
            if show_video:
                cv2.rectangle(orig, (x1, y1), (x2, y2), (0, 255, 0), 2)
                combined = np.hstack((orig, frame))
                cv2.imshow("Original | Cleaned", combined)
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break

        writer.release()
        if show_video:
            cv2.destroyAllWindows()
        self.cap.release()
        print(f"水印去除完成: {self.output_name}")

    def run(self):
        """
        1.视频去掉水印
        2.视频按时长切割成多个视频实现
        :return:
        """

        if self.logo_x1 is None:
            self._get_roi()
        # ---------------------------
        # 1. Step 1：去水印
        self.remove_watermark_ffmpeg()
        # self.remove_watermark(show_video=True)
        print(f"去水印完成：{self.output_name}")



# ---------------- 使用示例 ----------------
if __name__ == "__main__":
    Video_watermark_removal = VideoWatermarkRemoval(input_video_path="video.mp4")
    Video_watermark_removal.run()# 执行去水印
